import cv2

from .constants import *



def calc_back_proj_model(img, color_space, mask=None, num_of_bins=127, threshold=None):
    """求出用于 back projection 的模型(2D直方图)
    :param img: 已转换颜色空间的图片
    :param color_space: 图片所使用的颜色空间
    :param mask: 掩膜
    :param num_of_bins: 直方图区间的数量
    :param threshold: 二值化的阈值，为None则不对直方图做二值化
    """
    
    if color_space == SP_HSV:
        channels = [0, 1]  # using H, S channles
    elif color_space == SP_LAB:
        channels = [1, 2]  # using A, B channels
    else:
        raise ValueError("unsupported color space: {}".format(color_space))
    
    
    ranges = [0, 255, 0, 255]
    hist_size = [num_of_bins] * len(channels)
    
    hist = cv2.calcHist([img], channels, mask, hist_size, ranges, accumulate=False)
    hist = cv2.normalize(hist, hist, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX)
    
    if threshold is not None:
        _, hist = cv2.threshold(hist, threshold, maxval=255, type=cv2.THRESH_BINARY)
    
    return hist


def calc_back_proj(converted_img, color_space, hist):
    """
    包装了一下cv2.calcBackProject, 隐藏了一些无需理会的参数, 封装了与应用相关的操作,
    通过参数暴露出来.
    :param converted_img: 转换到HSV或LAB空间的图像
    :param color_space: 颜色空间, SP_HSV 或 SP_LAB
    :param hist: backproject直方图
    """
    
    if sp == SP_HSV:
        im_bkpj = cv2.calcBackProject([converted_img], [0,1], hist=hist, ranges=[0, 255]*2, scale=1)
    elif sp == SP_LAB:
        im_bkpj = cv2.calcBackProject([converted_img], [1,2], hist=hist, ranges=[0, 255]*2, scale=1)
    
    return im_bkpj


def png2hist(filename, color_space, num_of_bins=127, threshold=None):
    """读取png文件, alpha通道作为mask, 计算并返回mask选区的颜色直方图."""
    bgra = cv2.imread(filename, cv2.IMREAD_UNCHANGED)
    bgr  = bgra[:,:,:3]
    mask = bgra[:,:,3]
    
    cv2.imshow('bgr', bgr)
    cv2.imshow('mask', mask)
    if color_space == SP_HSV:
        img = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
    elif color_space == SP_LAB:
        img = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB)
    else:
        raise ValueError("unsupported color space: {}".format(color_space)) 
    
    hist = calc_back_proj_model(img, color_space, mask=mask, num_of_bins=num_of_bins, threshold=threshold)
    
    return hist


def load_model(filename):
    """包装成函数, 少写个参数."""
    hist = cv2.imread(filename, cv2.IMREAD_UNCHANGED)
    return hist
    

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    
    # im = cv2.imread(r'C:\_AAA\BlueNet\color_seg\a-template.png')
    # im = cv2.cvtColor(im, cv2.COLOR_BGR2LAB)
    # hist = calc_back_proj_model(im, SP_LAB, threshold=None)
    # print(hist.shape, hist.dtype, hist.max(), hist.min())
    
    # cv2.imshow('', hist)
    # cv2.waitKey(-1)
    
    # plt.imshow(hist)
    # plt.show()  
    
    # ===== ===== ===== ===== ===== =====
    # test png2hist()
    sp = SP_HSV
    
    hist = png2hist('../mask-lane2.png', color_space=sp, num_of_bins=100, threshold=10)
    cv2.imshow('hsit', hist)

    
    # ====== ===== ===== ===== ===== =====
    # explore the file storage
    # print('hist.shape: ', hist.shape)
    
    # hist_fname = '../model.png'
    # cv2.imwrite(hist_fname, hist)
    
    # hist = cv2.imread(hist_fname, cv2.IMREAD_UNCHANGED)
    # print("saved and read hist's shape:", hist.shape)
    
    # test in video
    cap = cv2.VideoCapture(r'../../../img_and_videos/lane-to-morning.mp4')
    
    while True:
        ret, raw = cap.read()
        raw = cv2.resize(raw, dsize=(0,0), fx=0.5, fy=0.5)  # 图小一些, 屏幕好放
        
        img = cv2.cvtColor(raw, cv2.COLOR_BGR2HSV) if sp==SP_HSV else cv2.cvtColor(raw, cv2.COLOR_BGR2LAB)
        im_bkpj = calc_back_proj(img, color_space=sp, hist=hist)
        im_lv = cv2.inRange(img[:,:,2], 50, 230)   # l of lab, or v of hsv
        
        im_bkpj = cv2.bitwise_and(im_bkpj, im_lv)
        # print(im_bkpj.shape)
        imtoshow = cv2.hconcat([raw, cv2.merge([im_bkpj]*3)])
        cv2.imshow('', imtoshow)
        # cv2.imshow('raw', raw)
        # cv2.imshow('bkpj', im_bkpj)
    
        k = cv2.waitKey(50) & 0xFF
        if k == 27 or not ret:  # Esc
            exit(0)
